#pragma once

#include <iostream>
#include <stack>

namespace solution_099
{

struct TreeNode
{
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

class Solution
{
public:
    void recoverTree(TreeNode* root)
    {
        if (root)
            traverse_mirrors(root);
    }

    void traverse(TreeNode* root)
    {
        std::stack<TreeNode*> stack;
        TreeNode* last_visited = nullptr;
        TreeNode* first = nullptr;
        TreeNode* second = nullptr;

        while (!stack.empty() || root) {
            while (root) {
                stack.push(root);
                root = root->left;
            }

            root = stack.top();
            stack.pop();
            if (last_visited && last_visited->val > root->val) {
                second = root;
                if (first)
                    break;
                first = last_visited;
            }
            last_visited = root;
            root = root->right;
        }

        std::swap(first->val, second->val);
    }

    void traverse_mirrors(TreeNode* root)
    {
        TreeNode* last_visited = nullptr;
        TreeNode* first = nullptr;
        TreeNode* second = nullptr;

        while (root) {
            if (root->left) {
                auto pred = root->left;
                while (pred->right && pred->right != root)
                    pred = pred->right;

                if (pred->right == nullptr) {
                    pred->right = root;
                    root = root->left;
                } else {
                    if (last_visited && last_visited->val > root->val) {
                        second = root;
                        if (first == nullptr)
                            first = last_visited;
                    }
                    last_visited = root;
                    root = root->right;

                    pred->right = nullptr;
                }
            } else {
                if (last_visited && last_visited->val > root->val) {
                    second = root;
                    if (first == nullptr)
                        first = last_visited;
                }
                last_visited = root;
                root = root->right;
            }
        }

        std::swap(first->val, second->val);
    }
};

}  /* namespace solution_099 */
